#include <torch/extension.h>
#include <iostream>
#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor. ")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous. ")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

void iou_gpu_launch(const int box1_num,
const float* box1,
const int box2_num,
const float* box2,
float* result);

int iou_gpu(torch::Tensor box1, torch::Tensor box2, torch::Tensor result){
    CHECK_INPUT(box1);
    CHECK_INPUT(box2);
    CHECK_INPUT(result);
    int box1_num = box1.size(0);
    int box2_num = box2.size(0);
    int box1_dim = box1.size(1);
    int box2_dim = box2.size(1);
    printf("box1_dim=%d\n",box1_dim);
    if(box1_dim!=4 || box2_dim!=4)
        return -1;
    printf("len(box1)=%d, len(box2)=%d\n",box1_num,box2_num);
    const float* box1_ptr = box1.data_ptr<float>();
    const float* box2_ptr = box2.data_ptr<float>();
    float* result_ptr = result.data_ptr<float>();
    iou_gpu_launch(box1_num,box1_ptr,box2_num,box2_ptr,result_ptr);
    return 0;
}


int iou_cplus(torch::Tensor box1, torch::Tensor box2, torch::Tensor result){
    int box1_num = box1.size(0);
    int box2_num = box2.size(0);
    int box1_dim = box1.size(1);
    int box2_dim = box2.size(1);
    printf("box1_dim=%d\n",box1_dim);
    if(box1_dim!=4 || box2_dim!=4)
        return -1;
    printf("len(box1)=%d, len(box2)=%d\n",box1_num,box2_num);
    float* box1_ptr = box1.data_ptr<float>();
    float* box2_ptr = box2.data_ptr<float>();
    float* result_ptr = result.data_ptr<float>();
    for(int i=0;i<box1_num;i++){
        for(int j=0;j<box2_num;j++){
            // x1 y1 x2 y2
            float* box1_single = box1_ptr+box1_dim*i;
            float* box2_single = box2_ptr+box2_dim*j;
            float overlap_width = std::min(*(box1_single+2),*(box2_single+2)) -std::max(*(box1_single+0),*(box2_single+0));
            float overlap_height = std::min(*(box1_single+3),*(box2_single+3)) -std::max(*(box1_single+1),*(box2_single+1));
            float overlap_area = overlap_width*overlap_height;
            float box1_area = (*(box1_single+2)-*(box1_single+0))*(*(box1_single+3)-*(box1_single+1));
            float box2_area = (*(box2_single+2)-*(box2_single+0))*(*(box2_single+3)-*(box2_single+1));
            if(box1_area<=0||box2_area<=0 ||overlap_width<=0||overlap_height<=0)
                *(result_ptr+i*box2_num+j)=0.0;
            else
                *(result_ptr+i*box2_num+j)=overlap_area/(box1_area+box2_area-overlap_area);
        }
    }
    return 0;
}


PYBIND11_MODULE( TORCH_EXTENSION_NAME, m )
{
    m.def("forward_cuda", &iou_gpu, "test");
    m.def("forward_cplus", &iou_cplus, "test2");
}